'''
Created on Oct 14, 2011

@author: u

Utility functions to extract pattern 
'''

import png
import numpy as np
from os import path
from hopfield import bin_to_hop, hop_to_bin

img_path = path.join(path.dirname(__file__), "images")

bit_depth = 8  #number of bits in gray scale

def read_png(filename):
    with open(path.join(img_path, filename), "rb") as pngfile:
        pngReader=png.Reader(file=pngfile)
        row_count, column_count, pngdata, meta = pngReader.asDirect() #@UnusedVariable
        assert meta['bitdepth'] == bit_depth
        assert meta['planes'] == 1

        #returns a 2d array 
        return np.vstack(map(np.int8, pngdata))
    
def expand_img(image):
    """
    transforms a nbit int array into n binary 2d arrays
    """
    
    def process_bit(i):
        mask = 2**i
        return (image & mask) >> i    
        
    return np.vstack(map(process_bit, range(bit_depth)))

def read_image_patterns(*images):
    images_1d, shape = convert_to_1d(*(map(read_png,  images)))
    arrays = np.array(tuple(expand_img(img) for img in images_1d))
    
    #transpose the list of patterns to have e pixel bits indices in first dimension
    #x : patterns
    #y : bits
    #z : pixels
    sx, sy, sz = arrays.shape
    dx, dy, dz = np.indices(arrays.shape)
    arrays = arrays.reshape(sy, sx, sz)
    arrays = arrays[dy, dx, dz]
    arrays = arrays.reshape(sy, sx, sz)
    arrays = bin_to_hop(arrays)
    return arrays, shape

def convert_to_1d(*images):
    """
    convert all input images to vectors
    and check they all have the same dimension
    """
    shape1 = images[0].shape
    
    def conv_img(img):
        if img.shape != shape1:
            raise ValueError("Not all images have the same dimensions")
        
        return img.reshape(-1)
    
    return np.vstack(map(conv_img, images)), shape1
    
def write_png(filename, image_2d):
    column_count, row_count = image_2d.shape
    with open(path.join(img_path, filename), 'wb') as pngfile:
        pngWriter = png.Writer(column_count, row_count,
                               greyscale=True,
                               alpha=False,
                               bitdepth=8)
    
        pngWriter.write(pngfile, image_2d)
        
def collapse_img(images):
    """
    reverse operation for expand_img
    """
    
    result = np.zeros(images[0].shape, dtype="uint8")
    
    for i, image in enumerate(images): 
        result |= hop_to_bin(image).astype("uint8") << i
    
    return result

def write_image_pattern(filenames, images, pattern_shape):
    sx, sy, sz = images.shape
    dx, dy, dz = np.indices(images.shape)
    images = hop_to_bin(images)
    images = images.reshape(sy, sx, sz)
    images = images[dy, dx, dz]
    images = images.reshape(sy, sx, sz)
    
    images = map(collapse_img, images)
    
    for image, filename in zip(images, filenames):
        write_png(filename, image.reshape(pattern_shape))
    
if __name__ == "__main__":
    patterns = read_image_patterns("lena_gray.png")
    
    
